from torch.utils.data import DataLoader, Dataset
import numpy as np
import random
from PIL import Image
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
import PIL.ImageOps
import pandas as pd
import os


class Dataset(Dataset):
    def __init__(self, datatxt, transform=None, should_invert=True):
        fh = open(datatxt, 'r')
        prefix = '/media/citaa/e8d7592e-674f-496d-a4f2-14ce0f3d8098/' \
                 'dataset/release_v0/images'
        labels = []
        clini = []
        derm = []
        for line in fh:
            line = line.rstrip()
            words = line.split(',')
            clini.append(os.path.join(prefix, words[0]))
            derm.append(os.path.join(prefix, words[1]))
            labels.append([int(words[2]), int(words[3]), int(words[4]), int(words[5]),
                          int(words[6]), int(words[7]), int(words[8]), int(words[9])])
        self.clini = clini
        self.derm = derm
        self.labels = labels
        self.transform = transform

    def __getitem__(self, index):

        clini_img = Image.open(self.clini[index]).convert('RGB')
        derm_img = Image.open(self.derm[index]).convert('RGB')
        labels = self.labels[index]

        if self.transform is not None:
            clini_img = self.transform(clini_img)
            derm_img = self.transform(derm_img)

        return clini_img, derm_img, labels

    def __len__(self):
        return len(self.clini)

